package com.raven.play.datalink.jdbc;

import akka.stream.Attributes;
import akka.stream.Outlet;
import akka.stream.Shape;
import akka.stream.SourceShape;
import akka.stream.stage.AbstractOutHandler;
import akka.stream.stage.GraphStage;
import akka.stream.stage.GraphStageLogic;
import lombok.extern.slf4j.Slf4j;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;

/**
 * JDBC 数据源
 *
 * @param <O>
 */
@Slf4j
public abstract class JdbcSourceGraphStage<O> extends GraphStage<SourceShape<O>> {

    private final Outlet<O> out = Outlet.create("JdbcSource.out");
    private final SourceShape<O> shape = new SourceShape<>(out);
    private final DataSource dataSource;
    private final String querySql;
    private final StatementBuilder<O> statementBuilder;
    private final ResultSetHandler<O> resultHandler;

    public JdbcSourceGraphStage(DataSource dataSource, String querySql, StatementBuilder<O> statementBuilder, ResultSetHandler<O> resultHandler) {
        this.dataSource = dataSource;
        this.querySql = querySql;
        this.statementBuilder = statementBuilder;
        this.resultHandler = resultHandler;
    }

    @Override
    public SourceShape<O> shape() {
        return shape;
    }

    @Override
    public GraphStageLogic createLogic(Attributes inheritedAttributes) throws Exception {
        return new JdbcSourceGraphStageLogic(shape);
    }


    class JdbcSourceGraphStageLogic extends GraphStageLogic {

        private Connection connection;
        private PreparedStatement preparedStatement;
        private ResultSet resultSet;

        public JdbcSourceGraphStageLogic(Shape shape) {
            super(shape);
            this.setHandler(out, new AbstractOutHandler() {
                @Override
                public void onPull() throws Exception {
                    // 如果仍然有数据
                    if (resultSet.next()) {
                        O record = resultHandler.parse(resultSet);
                        push(out, record);
                    } else {
                        // 如果没有更多数据，完成该阶段
                        completeStage();
                    }
                }
            });
        }



        @Override
        public void preStart() throws Exception {
            try {
                connection = dataSource.getConnection();
                preparedStatement = connection.prepareStatement(querySql, ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY);
                preparedStatement.setFetchSize(2000);
                preparedStatement.setFetchDirection(ResultSet.FETCH_FORWARD);

                if (statementBuilder != null) {
                    statementBuilder.assembly(preparedStatement, null);
                }

                resultSet = preparedStatement.executeQuery();
            } catch (SQLException e) {
                log.error(e.getMessage(), e);
                failStage(e);
            }
        }

        @Override
        public void postStop() throws Exception {
            try {
                if (resultSet != null) resultSet.close();
                if (preparedStatement != null) preparedStatement.close();
                if (connection != null) connection.close();
            } catch (SQLException e) {
                // 处理关闭资源时的异常
                log.error(e.getMessage(), e);
            }
        }

    }
}
